// Copyright (C) 2010 and onwards Harik Shazeer Labs
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Noam Shazeer and Georges Harik

// #include <pthread.h>
#ifndef _GNU_SOURCE 
#define _GNU_SOURCE
#endif

#include <unistd.h>
#include <sys/mman.h>
#include "allocator.h"
#include "jobqueue.h"

// This is used to keep track of the ThreadAllocator for a thread
GlobalAllocator G_ALLOC;

#ifdef OSX
const uint64 G_MAX_PHYSICAL_MEMORY = 1ull << 23;
const uint64 G_NUM_ALLOCATION_BLOCKS = 10;
const uint64 G_ARRAY_CACHE_SIZE = 1ull << 20;

#else
const uint64 G_MAX_PHYSICAL_MEMORY = 1ull << 32;
const uint64 G_NUM_ALLOCATION_BLOCKS = 100;
const uint64 G_ARRAY_CACHE_SIZE = 1ull << 27;
#endif

const uint64 G_ALLOCATION_BLOCK_SIZE = G_MAX_PHYSICAL_MEMORY;
const uint64 G_ALLOCATION_SIZE = 1ull << 17;
const uint64 G_PAGE_SIZE = sysconf(_SC_PAGESIZE) / sizeof(uint64);
const uint64 G_FOLLOW_Q_SIZE = 1ull << 14;

const bool G_CACHE_ARRAYS = true;

const uint64 kSavedSize = (uint64) 1e6;

uint64 G_FRONTIER_PUSHES = 0;
uint64 G_FRONTIER_POPS = 0;

void Frontier::AddObject(Object * o) {
  bool res = MakeReachable(o);
  if (res) {
    o->link__ = (uint64) head_;
    head_ = o;
    list_expansion_size_ += o->ArraySize();
    seen_memory_ += ObjectWordSize(o);
    // debugging __sync_fetch_and_add(&G_FRONTIER_PUSHES, 1);
  }
}

// Add the range of Objects into the linked list
void Frontier::ExpandRange(Object ** begin, Object ** end) {
  for (Object ** run = begin; run != end; run++) AddObject(*run);
}
void Frontier::Expand() {
  if (!head_) return;
  Object * head = head_;
  head_ = (Object *) head_->link__;
  // debugging __sync_fetch_and_add(&G_FRONTIER_POPS, 1);
  ExpandRange(head->ArrayData(), head->ArrayData() + head->ArraySize());
}


AllocatorDependant::AllocatorDependant() {
  if (this != &G_ALLOC)
    G_ALLOC.AddDependant(this);
}
AllocatorDependant::~AllocatorDependant() {}

AllocationBlock::AllocationBlock() {
  owner_ = NULL;
  data_ = next_ = limit_ = 
    (uint64 *)
    mmap(NULL, G_ALLOCATION_BLOCK_SIZE * sizeof(uint64), 
	 PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE | MAP_NORESERVE, 
	 -1, 0); // last two parameters ignored
  CHECK(data_ != MAP_FAILED);  
}
void AllocationBlock::Shrink() {
  next_ = new_next_;  
  G_ALLOC.SetBlockLimit(this, next_);
  uint64 size = next_ - data_;
  uint64 rounded_size = RoundUpToMultipleOf(size, G_PAGE_SIZE);
  munmap(data_ + rounded_size, 
	 (G_ALLOCATION_BLOCK_SIZE - rounded_size) * sizeof(uint64));
  void *result = 
    mmap(data_ + rounded_size, 
	 (G_ALLOCATION_BLOCK_SIZE - rounded_size) * sizeof(uint64), 
	 PROT_READ | PROT_WRITE, 
	 MAP_ANON | MAP_PRIVATE | MAP_NORESERVE | MAP_FIXED, 
	 -1, 0); // last two parameters ignored
  CHECK(result != MAP_FAILED);
}

AllocationBlock::~AllocationBlock() {
  munmap(data_, G_ALLOCATION_BLOCK_SIZE * sizeof(uint64));
}

GlobalAllocator::GlobalAllocator() {
  memory_allocated_ = 0;
  blocks_ = new AllocationBlock[G_NUM_ALLOCATION_BLOCKS];
  array_cache_ = new Object *[G_ARRAY_CACHE_SIZE];
  for (uint64 i=0; i<G_ARRAY_CACHE_SIZE; i++) array_cache_[i] = NULL;
  pthread_mutex_init(&lock_, NULL);
  AddDependant(this);
  next_block_ = 0;
}

// Any blocks assigned to the global allocator have limit_ = next_
void GlobalAllocator::RecallBlock(AllocationBlock *b) {
  if (!b) return;
  SetBlockLimit(b, b->next_);
  if (!b->owner_) return;
  b->owner_->block_ = NULL;
  b->owner_ = NULL;
}

void GlobalAllocator::Allocate(ThreadAllocator *ta, uint64 needed) {
  // cout << "GlobalAllocator::AllocateBlock" << endl;
  ta->ReleaseLock();
  Lock(lock_, "::GlobalAllocator global_lock_");
  
  if (ta->block_) RecallBlock(ta->block_);
  do { 
    next_block_ = (next_block_ + 1) % G_NUM_ALLOCATION_BLOCKS;
  } while (!BlockIsUsable(&blocks_[next_block_]));
  
  uint64 allocation = max(needed, G_ALLOCATION_SIZE);
  if (memory_allocated_ + allocation > G_MAX_PHYSICAL_MEMORY) {
    GarbageCollect(allocation);
  }
  AllocationBlock * b = &blocks_[next_block_];
  b->owner_ = ta;
  ta->block_ = b;
  CHECK(b->next_ == b->limit_);
  SetBlockLimit(b, b->limit_ + allocation);
  
  // Release the global allocator lock
  ta->AcquireLock();
  Unlock(lock_, "::GetBlock gc_lock_");
  ta->CheckException();

  //cout << "end GlobalAllocator::AllocateBlock" << endl;
}

ThreadAllocator::ThreadAllocator()
 :   saved_(kSavedSize)
{
  array_cache_tries_ = array_cache_hits_ = 0;
  G_ALLOC.AddThreadAllocator(this);
  pthread_mutex_init(&lock_, NULL);
  block_ = NULL;
  throw_exception_ = false;
}

uint64 * ThreadAllocator::Allocate(uint64 length) {
  if (!block_ || block_->AvailableSize() < length) 
    G_ALLOC.Allocate(this, length);
  CHECK(block_->AvailableSize() >= length);
  uint64 * ret = block_->next_;
  block_->next_ += length;
  return ret;
}

Object * ThreadAllocator::AllocateArray(uint64 length) {
  Object * object = 
    (Object *)Allocate(sizeof(Object) / sizeof(Object *) + length);
  
  object->hash_ = (uint64) object;
  object->object_info_ = 0;
  object->garbage_ = 0;
  object->SetType(ARRAY_TYPE);
  object->SetFlag(MUTABLE_OBJECT); 
  object->SetLength(length);
  JobQueueZero(G_JQ, object->ArrayData(), length * sizeof(Object *));
  // for (uint64 c = 0; c<length; c++) object->Ref(c) = NULL;
  return object;
}

Object * ThreadAllocator::MakeArray(Object ** contents, uint64 length){
  uint64 h = ArrayHash(contents, length);
  if (G_CACHE_ARRAYS) {
    Object * o = G_ALLOC.array_cache_[h % G_ARRAY_CACHE_SIZE];
    /*if (array_cache_tries_ % 1000000 == 0) 
      cout << "Array cache hit rate: " 
      << ( (double)array_cache_hits_ / array_cache_tries_) << endl;*/
    // array_cache_tries_++;
    if (o && (o->hash_ == h)) {
      // array_cache_hits_++;
      return o;
    }
  }
  Object * ret = AllocateArray(length);
  ret->MakeImmutable();
  ret->hash_ = h;
  memcpy(ret->ArrayData(), contents, length * sizeof(Object *));
  ret->ComputeNodeCount();
  if (G_CACHE_ARRAYS) G_ALLOC.array_cache_[h % G_ARRAY_CACHE_SIZE] = ret;
  return ret;
}

// Access ToString for debugging
extern string ToString(Object * o);

// ---------- Garbage Collection from here on ----------



struct ExpandFrontierTask : public JobQueue::Task {
  ExpandFrontierTask(Frontier frontier) : frontier_(frontier) {}
  void Run() {
    // Add all objects in a range to the linked list
    frontier_.ExpandRange(frontier_.start_, frontier_.end_); 
    while (frontier_.head_) {
      SplitWhileNecessary();
      frontier_.Expand();
    }
    __sync_fetch_and_add(&G_ALLOC.memory_used_, frontier_.seen_memory_);
    // cout << "Frontier task done" << endl;
  }
  // For load balancing, we recommend implementing this
  void SplitWhileNecessary() {}
  Frontier frontier_;
};

// Insert reachable objects for a shard of an AllocatorDependant
// into a frontier and push it onto the global allocator
struct FlagReferencesTask : public JobQueue::Task {
  FlagReferencesTask (AllocatorDependant * ad, 
		      int32 shard)
    : ad_(ad), shard_(shard) {}
  void Run() { 
    ad_->FlagReferences(shard_, &frontier_);
    G_ALLOC.propagate_job_->AddTask
      (new ExpandFrontierTask(frontier_));
  }
  AllocatorDependant * ad_;
  int32 shard_;
  Frontier frontier_;
};


void GlobalAllocator::ComputeReachable() {
  // Help keep track of memory used
  memory_used_ = 0;

  propagate_job_ = new JobQueue::Job(G_JQ);

  // Create the initial frontiers
  cout << "::FlagReferences" << endl;
  JobQueue::Job j(G_JQ);
  forall (run, dependants_) {
    AllocatorDependant * ad = *run;
    int32 num_shards = ad->NumShards();
    for (int32 c=0; c<num_shards; c++)
      j.AddTask(new FlagReferencesTask(ad, c));
  }
  j.Wait();
  
  // Now expand the frontiers
  propagate_job_->Wait();
  delete propagate_job_;
}

struct ZeroReachabilityTask : public BlockTask {
  void Handle(Object * o) { o->RemoveFlag(REACHABLE_OBJECT); } 
};

struct VerifyReachabilityTask : public BlockTask {
  void Handle(Object * o) {
    if (o->IsReachable()) {
      for (uint64 i=0; i<o->ArraySize(); i++) {
	Object * o2 = o->Ref(i);
	if (o2 && IsBoxed(o2)) {
	  CHECK(GetType(o2) == ARRAY_TYPE);
	  CHECK(o2->IsReachable());
	}
      }
    }
  } 
};

struct FindNewLocationTask : public BlockTask {
  void StartBlock() { write_pointer_ = b_->data_; }
  void Handle(Object *o) {
    if (o->IsReachable()) {
      o->new_location__ = (uint64) write_pointer_;
      write_pointer_ += ObjectWordSize(o);
    } else ReplaceWithFreeSpace(o);
  }
  void FinishBlock() {
    b_->new_next_ = write_pointer_;
  }
  uint64 * write_pointer_; 
};

struct RewriteArraysTask : public BlockTask {
  void Handle(Object *o) {
    if (GetType(o) != ARRAY_TYPE) return;
    for (uint64 i=0; i<o->ArraySize(); i++) {
      Object * old = o->Ref(i);
      Object * nu = NewLocation(o->Ref(i));
      if (old && !nu) {
	cout << "Rewriting array reference to null" << endl;
	CHECK(false);
      }
      o->Ref(i) = NewLocation(o->Ref(i));	
    }
  }
};

struct MoveObjectsTask : public BlockTask {
  void Handle(Object *o) {
    if (!WasDeleted(o)) {
      void * new_location = NewLocation(o);
      memmove(new_location, o, sizeof(uint64) * ObjectWordSize(o));
      Object * new_o = (Object *) new_location;
      new_o->garbage_ = 0;
      new_o->RemoveFlag(REACHABLE_OBJECT);
    }
  }
};


// this function has to get and release locks
void GlobalAllocator::GarbageCollect(uint64 nbytes) {
  cout << "::GarbageCollect needs nbytes:" << nbytes << endl;
  cout << "memory_allocated_: " << memory_allocated_ << endl;

  // Acquire all the run locks
  for (uint32 c=0; c<thread_allocators_.size(); c++) {
    ThreadAllocator * ta = thread_allocators_[c];
    ta->AcquireLock();
    RecallBlock(ta->block_);
  }

  uint64 urgency = 0;
  while (1) {
    if (urgency == KILL_PROCESS_URGENCY) {
      cout << "Garbage collect gives up at urgency " << KILL_PROCESS_URGENCY
	   << ".  Kill process" << endl;
      exit(1);
    }
    G_ALLOC.urgency_ = urgency;
    //     RipAllObjects<VerifyReachabilityTask>();
    ComputeReachable();
    //    RipAllObjects<VerifyReachabilityTask>();
    cout << "Urgency " << urgency << 
      " words used: " << memory_used_ << endl;
    if (HasEnoughMemory(nbytes)) break;
    urgency++;
    RipAllObjects<ZeroReachabilityTask>();
  }

  // Determine the new location for all objects
  cout << "FindNewLocationTask" << endl;
  RipAllObjects<FindNewLocationTask>();

  // Notify external Object* holders 
  // their pointers are about to become invalid
  NotifyObjectsDeleted();

  // Rewrite all internal Object*s in objects to be kept
  cout << "::RewriteArraysTask" << endl;
  RipAllObjects<RewriteArraysTask>();
 
  // Actually move the objects
  cout << "::MoveObjectsTask" << endl;
  RipAllObjects<MoveObjectsTask>();

  // Shrink all the blocks
  cout << "::ShrinkBlocks" << endl;
  for (uint32 c=0; c< G_NUM_ALLOCATION_BLOCKS; c++) {
    G_ALLOC.blocks_[c].Shrink();
  }
  cout << "GarbageCollect:: " << endl;

  // Release everyone
  for (uint32 c=0; c<thread_allocators_.size(); c++) {
    ThreadAllocator * ta = thread_allocators_[c];
    // if (!threads_alive_) ta->throw_exception_ = true;
    ta->ReleaseLock();
  }
}

bool GlobalAllocator::HasEnoughMemory(uint64 nbytes) {
  cout << "memory_used_:" << memory_used_ << endl;
  return (memory_used_ * 2 + nbytes < G_MAX_PHYSICAL_MEMORY);
}

struct ObjectsDeletedTask : public JobQueue::Task {
  ObjectsDeletedTask (AllocatorDependant * ad, int32 shard) :
    ad_(ad), shard_(shard) {}
  void Run() { ad_->ObjectsDeleted(shard_); }
  AllocatorDependant * ad_;
  int32 shard_;
};

// The ObjectsDeleted call on an allocatorDependant lets it know
// that any Object *s it owns are about to be moved or deleted,
// so it may want to update its external Object *s before that happens
// or react in some way
void GlobalAllocator::NotifyObjectsDeleted() {
  cout << "::NotifyObjectsDeleted" << endl;
  JobQueue::Job j(G_JQ);
  forall (run, dependants_) {
    AllocatorDependant * ad = *run;
    int32 num_shards = ad->NumShards();
    for (int32 c=0; c<num_shards; c++) {
      j.AddTask(new ObjectsDeletedTask(ad, c));
    }
  }
  j.Wait();
}

void GlobalAllocator::FlagReferences(int32 shard, Frontier * frontier) {}

void GlobalAllocator::ObjectsDeleted(int32 shard) {
  if (!G_CACHE_ARRAYS) return;
  uint64 begin = ShardBegin(G_ARRAY_CACHE_SIZE, NumShards(), shard);
  uint64 end = ShardEnd(G_ARRAY_CACHE_SIZE, NumShards(), shard);
  for (uint64 i=begin; i<end; i++) {
    array_cache_[i] = NewLocation(array_cache_[i]);
  }
}

// Mark what you know about as externally important
void ThreadAllocator::FlagReferences(int32 shard, Frontier * frontier) {
  if (G_ALLOC.urgency_ >= KILL_THREADS_URGENCY) {
    throw_exception_ = true;
    return;
  }
  CHECK(false);
  cout << "ThreadAllocator::FlagReferences" << endl;

  for (int64 c=0; c<saved_.size(); c++) {
    Object * o = (*(saved_[c]));
    // cout << "o:" << o << endl;
    // cout << "o:" << ToString(o) << endl;    
    frontier->AddObject(o);
  }
  // cout << "vectors now" << endl;
  for (uint64 c=0; c<saved_vectors_.size(); c++) {
    vector<Object *> * v = saved_vectors_[c];
    for (uint64 c2=0; c2<v->size(); c2++) {
      Object * o = (*v)[c2];
      // cout << "o:" << o << endl;
      frontier->AddObject(o);
    }
  }
  cout << "end ThreadAllocator::FlagReferences" << endl;
}

void ThreadAllocator::ObjectsDeleted(int32 shard) {
  if (G_ALLOC.urgency_ >= KILL_THREADS_URGENCY) return;
  // cout << "ThreadAllocator::ObjectsDeleted" << endl;

  for (int64 c=0; c<saved_.size(); c++) {
    Object * o = (*(saved_[c]));
    CHECK(!WasDeleted(o));
    (*(saved_[c])) = NewLocation(o);
  }
  // cout << "vectors now" << endl;
  for (uint64 c=0; c<saved_vectors_.size(); c++) {
    vector<Object *> * v = saved_vectors_[c];
    for (uint64 c2=0; c2<v->size(); c2++) {
      Object * o = (*v)[c2];
      CHECK(!WasDeleted(o));
      (*v)[c2] = NewLocation(o);
    }
  }
  // cout << "end ThreadAllocator::ObjectsDeleted" << endl;
}

/*

struct DeleteLowPriorityTask : public BlockTask {
  void Handle(Object *o) {
    if (Priority(o) < G_ALLOC.priority_threshold_) 
      o->SetFlag(NO_EXTERNAL_IMPORTANCE);
  }
};

struct ClearImportanceFlagTask : public BlockTask {
  void Handle (Object *o) {
    o->RemoveFlag(NO_EXTERNAL_IMPORTANCE);
  }
};

*/

void VectorObjectClear(Object *& v) { 
  if (!v) return;
  CHECK(GetType(v) == ARRAY_TYPE);
  CHECK(v->IsMutable());
  v = NULL;
}
uint64 VectorObjectSize(Object * v) {
  if (!v) return 0;
  return GetInteger(v->Ref(0));
}
void VectorObjectSetSize(Object * v, uint64 size) {
  CHECK(v);
  CHECK(size <= v->ArraySize() - 1);
  v->Ref(0) = MakeInteger(size);
}

Object *& VectorObjectRef(Object *v, uint64 pos) {
  CHECK(GetType(v) == ARRAY_TYPE);
  CHECK(v->IsMutable());
  // CHECK(v && (pos < VectorObjectSize(v)));
  return v->Ref(pos+1);
}

// pushes increment NULLs onto the end of the VectorObject.  
// If the VectorObject is garbage collected, returns false.
// This is thread unsafe. You need to make sure only 1 thread runs this.
bool VectorObjectExtend(ThreadAllocator *a, 
			       Object *& v, 
			       uint64 increment) {
  if (v) {
    CHECK(GetType(v) == ARRAY_TYPE);
    CHECK(v->IsMutable());
  }
  if (!v) {
    v = a->AllocateArray(increment + 1);
    v->Ref(0) = MakeInteger(0);
  }
  uint64 old_size = VectorObjectSize(v);
  if (old_size + increment + 1 > v->ArraySize()) {
    Object * new_array = a->AllocateArray(2 * (old_size + increment));
    if (v) { // need to check because it could have been garbage collected
      Object ** src = &v->Ref(0);
      Object ** dst = &new_array->Ref(0);
      JobQueueCopy(G_JQ, (void*)dst, (void*)src, 
		   sizeof(Object *) * (1 + old_size));
      v = new_array;
    } else {
      v = NULL;
      return false;
    }
  }
  VectorObjectSetSize(v, old_size + increment);
  return true;
}

// If the vector you are pointing to gets garbage collected during
// resizing to add an element, it just adds the element to an empty vector
// This is thread unsafe. You need to make sure only 1 thread runs this.
void VectorObjectPushBack(ThreadAllocator *a, 
				 Object *& v,
				 Object ** val,
				 uint64 length,
				 uint64 num_copies) {
  uint64 place = VectorObjectSize(v);
  uint64 increment = length * num_copies;
  if (!VectorObjectExtend(a, v, increment)) return;
  for (uint64 obj_num = 0; obj_num < length; obj_num++) {
    for (uint64 copy = 0; copy < num_copies; copy++) {
      v->Ref(place + 1) = *val;
      place++;
    }
    val++;
  }
}

void VectorObjectParallelPush
(ThreadAllocator *a,
 Object *& v,
 const vector<pair<pair<uint64, Object *> *, 
 pair<uint64, Object *> *> > & ranges) {
  uint64 increment = 0;
  for (uint64 i=0; i<ranges.size(); i++) 
    increment += ranges[i].second - ranges[i].first;
  uint64 place = VectorObjectSize(v);
  if (!VectorObjectExtend(a, v, increment)) return;
  JobQueue::Job j(G_JQ);
  for (uint64 i=0; i<ranges.size(); i++) {
    j.AddTask(new CopySecondTask<uint64, Object *>
	      (&VectorObjectRef(v, place), 
	       ranges[i].first, ranges[i].second));    
    place += ranges[i].second - ranges[i].first;
  }
  j.Wait();  
}
